/* Include files */
#include "math.h"
#include "mex.h"

/* Include to sort : http://stackoverflow.com/questions/1787996/c-library-function-to-do-sort */
/*
#include <stdio.h>
#include <stdlib.h>
*/

/*************************************************/
/* Delete Zero Vec */
/*
static int AssignNonZeroVec(int *rowIn, int *rowOut, int A)
{
	int a = 0; int b = 0;
	for (a = 0; a < A; a++)
	{
		if (rowIn[a] > 0) 
		{
			rowOut[b] = rowIn[a]; 
			b++;
			printf("%d b \n",b);
		}
		printf("%d A \n",a);
	}
	return b;
}
*/

/* Checks whether inputs are matrices */
#define IS_REAL_2D_FULL_DOUBLE(P) (!mxIsComplex(P) && \
mxGetNumberOfDimensions(P) == 2 && !mxIsSparse(P) && mxIsDouble(P))
#define IS_REAL_SCALAR(P) (IS_REAL_2D_FULL_DOUBLE(P) && mxGetNumberOfElements(P) == 1)

/*************************************************/
/* MAIN                                          */
/*************************************************/

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
	/* Define output matrix */
	#define Amean   		plhs[0]
	#define Bmean    		plhs[1]

	/* Define input matrices */
	#define Amat 		prhs[0]
	#define Bmat    	prhs[1]

	#define FirmMatch 	prhs[2]
	#define Jm          prhs[3]
	#define GroupIndex  prhs[4]
	#define GroupCount  prhs[5]

	#define MomentSwitch  prhs[6]

	/* Check the number of inputs is correct (Aoffer,Baccept,QbarA) */
	if(nrhs < 7)
		mexErrMsgTxt("Too few input arguments.");
	else if(nrhs > 7)
		mexErrMsgTxt("Too many input arguments.");

	if(!IS_REAL_2D_FULL_DOUBLE(Amat))
		mexErrMsgTxt("Amat must be real 2D full double array.");
	if(!IS_REAL_2D_FULL_DOUBLE(Bmat))
		mexErrMsgTxt("Bmat must be real 2D full double array.");
	
	if(!IS_REAL_2D_FULL_DOUBLE(FirmMatch))
		mexErrMsgTxt("FirmMatch must be real 2D full double array.");
	if(!IS_REAL_2D_FULL_DOUBLE(Jm))
		mexErrMsgTxt("Jm must be real 2D full double array.");
	if(!IS_REAL_2D_FULL_DOUBLE(GroupIndex))
		mexErrMsgTxt("GroupIndex must be real 2D full double array.");
	if(!IS_REAL_2D_FULL_DOUBLE(GroupCount))
		mexErrMsgTxt("GroupCount must be real 2D full double array.");
	if(!IS_REAL_2D_FULL_DOUBLE(MomentSwitch))
		mexErrMsgTxt("MomentSwitch must be real 2D full double array.");

	/* Dimensions and pointers for inputs */

	int N = mxGetM(Amat);
	int M = mxGetN(GroupIndex);
	int Maxlen = mxGetM(GroupIndex);

	int A = mxGetN(Amat);
	int B = mxGetN(Bmat);

	double *A_pt 		= mxGetPr(Amat);
	double *B_pt 		= mxGetPr(Bmat);
	double *M_pt 		= mxGetPr(MomentSwitch);

	double *Firm_pt 	= mxGetPr(FirmMatch);
	double *Jm_pt 		= mxGetPr(Jm);
	double *GpIn_pt		= mxGetPr(GroupIndex);
	double *GpCt_pt		= mxGetPr(GroupCount);

	/* Statistic containers */

	Amean = mxCreateDoubleMatrix(M,A,mxREAL);
	double *Amean_pt = mxGetPr(Amean);

	Bmean = mxCreateDoubleMatrix(M,B,mxREAL);
	double *Bmean_pt = mxGetPr(Bmean);

	int J = Jm_pt[0];
	int JM = Jm_pt[0]*M;
	
	/* Temporary containers */

	int *rowVec = calloc(Maxlen, sizeof (int));
	int *row_pt = rowVec;

	int *firmCt = calloc(J, sizeof (int));
	int *firmCt_pt = firmCt;

	/* ---------------------------------------------------------------------------- */
	/* MAIN                                                                         */
	/* ---------------------------------------------------------------------------- */

	int a, b, m, i, j;
	int len, marLen, index, marIndex, loc, locJ, totFirm;
	double temp1, temp2;

	int begin_pt = 0;

	/* ---------------------------------------------------------------------------- */
	/* Assignment Code - Within Market Means                                        */
	/* ---------------------------------------------------------------------------- */
	for (m = 0; m < M; m++)
	{
		/* ----------------------------------- */
		/* Delete zeros from MarketIndex vector */

		marLen = 0;
		for (len = 0; len < Maxlen; len++)
		{
			if (GpIn_pt[len + Maxlen*m] != 0)
			{
				row_pt[len] = GpIn_pt[len + Maxlen*m];
				marLen++;
			}
		}

		/* ----------------------------------- */
		/* Mean Code */

		for(len = 0; len < marLen; len++)
		{
			marIndex = row_pt[len] - 1;
			if(Firm_pt[marIndex] > 0)
			{
				for(a = 0; a < A; a++)
					Amean_pt[m + M*a] = Amean_pt[m + M*a] + A_pt[marIndex + N*a]/(float)marLen;
				for(b = 0; b < B; b++)
					Bmean_pt[m + M*b] = Bmean_pt[m + M*b] + B_pt[marIndex + N*b]/(float)marLen;
			}
		}
	}

	/* ---------------------------------------------------------------------------- */
	/* Match-level Moments: Joint distribution and covariance of matched attributes */
	/* ---------------------------------------------------------------------------- */

	if (M_pt[0] == 1)
	{

		#define JntDist 		plhs[2]
		#define Cov 			plhs[3]

		JntDist = mxCreateDoubleMatrix(1,A*B,mxREAL);
		double *JntDist_pt = mxGetPr(JntDist);

		Cov = mxCreateDoubleMatrix(1,A*B,mxREAL);
		double *Cov_pt = mxGetPr(Cov);

		double amean_temp = 0;
		double bmean_temp = 0;
		index = 0;

		for (a = 0; a < A; a++)
		{
			for (b = 0; b < B; b++)
			{
				for (i = 0; i < N; i++)
				{
					amean_temp = amean_temp + A_pt[i + N*a]/(float)N;
					bmean_temp = bmean_temp + B_pt[i + N*b]/(float)N;
				}
				for (i = 0; i < N; i++)
				{
					JntDist_pt[index] = JntDist_pt[index] + A_pt[i + N*a]*B_pt[i + N*b]/(float)N;
					Cov_pt[index] = Cov_pt[index] + (A_pt[i + N*a] - amean_temp)*(B_pt[i + N*b] - bmean_temp)/(float)(N - 1);
				}

				amean_temp = 0;
				bmean_temp = 0;
				index++;
			}
		}
	}

	/* ---------------------------------------------------------------------------- */
	/* Within-Market Moments: Variance and covariance of firm, parcel att.          */
	/* ---------------------------------------------------------------------------- */

	if (M_pt[1] == 1)
	{
		#define Avar   			plhs[4]
		#define Bvar    		plhs[5]
		#define ABcovar   		plhs[6]

		Avar = mxCreateDoubleMatrix(M,A,mxREAL);
		double *Avar_pt = mxGetPr(Avar);

		Bvar = mxCreateDoubleMatrix(M,B,mxREAL);
		double *Bvar_pt = mxGetPr(Bvar);

		ABcovar = mxCreateDoubleMatrix(M,A*B,mxREAL);
		double *ABcovar_pt = mxGetPr(ABcovar);


		for (m = 0; m < M; m++)
		{
			/* ----------------------------------- */
			/* Delete zeros */

			marLen = 0;
			for (len = 0; len < Maxlen; len++)
			{
				if (GpIn_pt[len + Maxlen*m] != 0)
				{
					row_pt[len] = GpIn_pt[len + Maxlen*m];
					marLen++;
				}
			}

			/* ----------------------------------- */
			/* Calculate Moments */

			if (marLen > 1)
			{
				for(len = 0; len < marLen; len++)
				{
					marIndex = row_pt[len] - 1;
					index = 0;

					/* Calculate variance */
					for (a = 0; a < A; a++)
					{
						temp1 			 = A_pt[marIndex + N*a] - Amean_pt[m + M*a];
						Avar_pt[m + M*a] = Avar_pt[m + M*a] + (temp1*temp1)/(float)(marLen - 1);
					}
					for (b = 0; b < B; b++)
					{
						temp1 			 = B_pt[marIndex + N*b] - Bmean_pt[m + M*b];
						Bvar_pt[m + M*b] = Bvar_pt[m + M*b] + (temp1*temp1)/(float)(marLen - 1);
					}

					/* Calculate covariance */
					for (b = 0; b < B; b++)
					{
						temp2 = B_pt[marIndex + N*b] - Bmean_pt[m + M*b];
						for (a = 0; a < A; a++)
						{
							temp1 = A_pt[marIndex + N*a] - Amean_pt[m + M*a];
							ABcovar_pt[m + M*index] = ABcovar_pt[m + M*index] + (temp1*temp2)/(float)(marLen - 1);
							index++;
						}
					}
				}
			}
		}
	}

	/* ---------------------------------------------------------------------------- */
	/* Within-Firm Moments: Mean, Var of matched attributes grouped by firm         */
	/* ---------------------------------------------------------------------------- */

	if (M_pt[2] == 1)
	{
		#define Amean_firm		plhs[7]
		#define Avar_firm		plhs[8]

		Amean_firm = mxCreateDoubleMatrix(J,A,mxREAL);
		double *Amean_firm_pt = mxGetPr(Amean_firm);

		Avar_firm = mxCreateDoubleMatrix(J,A,mxREAL);
		double *Avar_firm_pt = mxGetPr(Avar_firm);

		/* Temporary Containers */
		int    *firmSign = calloc(J, sizeof (int));
		int    *firmS_pt = firmSign;

		int    FirmID = 0;

		/* Initialize to zero */
		for (j = 0; j < J; j++)
			firmS_pt[j] = 0;

		/* Count of leases signed by a firm */
		for (i = 0; i < N; i++)
		{
			FirmID = Firm_pt[i] - 1;
			if (FirmID >= 0)
				firmS_pt[FirmID]  = firmS_pt[FirmID] + 1;
		}

		/* Mean and variance of leases signed by a firm */
		for (a = 0; a < A; a++)
		{
			for (i = 0; i < N; i++)
			{
				FirmID = Firm_pt[i] - 1;
				if (firmS_pt[FirmID] > 0 & FirmID >= 0)
					Amean_firm_pt[FirmID + J*a] = Amean_firm_pt[FirmID + J*a] + A_pt[i + N*a]/(float)(firmS_pt[FirmID]);
			}
			for (i = 0; i < N; i++)
			{
				FirmID = Firm_pt[i] - 1;
				if (firmS_pt[FirmID] > 1 & FirmID >= 0)
					Avar_firm_pt[FirmID + J*a]  = Avar_firm_pt[FirmID + J*a]  + (A_pt[i + N*a] - Amean_firm_pt[FirmID + J*a])*(A_pt[i + N*a] - Amean_firm_pt[FirmID + J*a])/(float)(firmS_pt[FirmID] - 1);
			}
		}

		free(firmSign);
	}

	/* ---------------------------------------------------------------------------- */
	/* Within-Firm Moments: Joint distribution and covariance of firm, parcel att.  */
	/* ---------------------------------------------------------------------------- */

	if (M_pt[3] == 1)
	{

		#define ABdense_firm		plhs[9]
		#define ABcovar_firm		plhs[10]

		ABdense_firm = mxCreateDoubleMatrix(J,A*B,mxREAL);
		double *ABdense_firm_pt = mxGetPr(ABdense_firm);

		ABcovar_firm = mxCreateDoubleMatrix(J,A*B,mxREAL);
		double *ABcovar_firm_pt = mxGetPr(ABcovar_firm);

		/* Temporary Containers */
		int    *firmSign = calloc(J, sizeof (int));
		int    *firmS_pt = firmSign;

		double *Atemp = calloc(J*A, sizeof (double));
		double *Atemp_pt = Atemp;

		double *Btemp = calloc(J*B, sizeof (double));
		double *Btemp_pt = Btemp;

		int    FirmID = 0;

		/* Initialize to zero */
		for (j = 0; j < J; j++)
			firmS_pt[j] = 0;

		/* Count of leases signed by a firm */
		for (i = 0; i < N; i++)
		{
			FirmID = Firm_pt[i] - 1;
			if (FirmID >= 0)
				firmS_pt[FirmID]  = firmS_pt[FirmID] + 1;
		}

		/* Mean and covariance of leases signed by a firm */
		for (a = 0; a < A; a++)
		{
			for (i = 0; i < N; i++)
			{
				FirmID = Firm_pt[i] - 1;
				if (firmS_pt[FirmID] > 0 & FirmID >= 0)
					Atemp_pt[FirmID + J*a] = Atemp_pt[FirmID + J*a] + A_pt[i + N*a]/(float)(firmS_pt[FirmID]);
			}
		}
		for (b = 0; b < B; b++)
		{
			for (i = 0; i < N; i++)
			{
				FirmID = Firm_pt[i] - 1;
				if (firmS_pt[FirmID] > 0 & FirmID >= 0)
					Btemp_pt[FirmID + J*b] = Btemp_pt[FirmID + J*b] + B_pt[i + N*b]/(float)(firmS_pt[FirmID]);
			}
		}

		index = 0;
		for (a = 0; a < A; a++)
		{
			for (b = 0; b < B; b++)
			{
				for (i = 0; i < N; i++)
				{
					FirmID = Firm_pt[i] - 1;
					if (firmS_pt[FirmID] > 1 & FirmID >= 0)
						ABcovar_firm_pt[FirmID + J*index] = ABcovar_firm_pt[FirmID + J*index] + (A_pt[i + N*a] - Atemp_pt[FirmID + J*a])*(B_pt[i + N*b] - Btemp_pt[FirmID + J*b])/(float)(firmS_pt[FirmID] - 1);
					if (firmS_pt[FirmID] > 0 & FirmID >= 0)
						ABdense_firm_pt[FirmID + J*index] = ABdense_firm_pt[FirmID + J*index] + (A_pt[i + N*a]*B_pt[i + N*b])/(float)(firmS_pt[FirmID]);
				}
				index++;
			}
		}

		free(Atemp);
		free(Btemp);
		free(firmSign);
	}

	/* ---------------------------------------------------------------------------- */
	/* Within-Market-Firm Moments: Variance and covariance of firm, parcel att.     */
	/* ---------------------------------------------------------------------------- */

	if (M_pt[4] == 1)
	{

		#define GroupMean_mean 	plhs[11]
		#define GroupVar_mean  	plhs[12]

		#define GroupMean   	plhs[13]
		#define GroupVar   		plhs[14]

		GroupMean = mxCreateDoubleMatrix(JM,A,mxREAL);
		double *GpMean_pt = mxGetPr(GroupMean);

		GroupVar = mxCreateDoubleMatrix(JM,A,mxREAL);
		double *GpVar_pt = mxGetPr(GroupVar);

		GroupMean_mean = mxCreateDoubleMatrix(M,A,mxREAL);
		double *GpMean_mean_pt = mxGetPr(GroupMean_mean);

		GroupVar_mean = mxCreateDoubleMatrix(M,A,mxREAL);
		double *GpVar_mean_pt = mxGetPr(GroupVar_mean);

		for (m = 0; m < M; m++)
		{
			/* ----------------------------------- */
			/* Initialize to zero */
			
			for (j = 0; j < J; j++)
				firmCt_pt[j] = 0;

			/* ----------------------------------- */
			/* Delete zeros from MarketIndex vector */

			marLen = 0;
			for (len = 0; len < Maxlen; len++)
			{
				if (GpIn_pt[len + Maxlen*m] != 0)
				{
					row_pt[len] = GpIn_pt[len + Maxlen*m];
					marLen++;
				}
			}

			if (GpCt_pt[m + M*0] > 1)
			{
				/* ----------------------------------- */
				/* Calculate Group Moments */

				/* Count the number of leases signed in a market m by firm j */
				for(len = 0; len < marLen; len++)						/* Loop through all observations in market m */
				{
					loc 		= Firm_pt[len] - 1;						/* Firm ID */
					if (loc > -1)
						firmCt_pt[loc] = firmCt_pt[loc] + 1;
				}

				/* Count the number of active firms in market m */
				totFirm = 0;
				for (j = 0; j < J; j++)
				{
					if (firmCt_pt[j] > 0)
						totFirm++;
				}

				/* Firm-level mean for a market */
				for(len = 0; len < marLen; len++)						/* Loop through all observations in market m */
				{
					marIndex 	= row_pt[len] - 1;
					loc 		= Firm_pt[marIndex] - 1;				/* Firm ID */
					if (loc > -1 & firmCt_pt[loc] > 0)
					{
						for (a = 0; a < A; a++)
							GpMean_pt[J*m + loc + JM*a] = GpMean_pt[J*m + loc + JM*a] + A_pt[marIndex + N*a]/(float)firmCt_pt[loc];
					}
				}

				/* Firm-level var for a market */
				for(len = 0; len < marLen; len++)						/* Loop through all observations in market m */
				{
					marIndex 	= row_pt[len] - 1;
					loc 		= Firm_pt[marIndex] - 1;				/* Firm ID */
					if (loc > -1 & firmCt_pt[loc] > 1)
					{
						for (a = 0; a < A; a++)
						{
							temp1 					   = A_pt[marIndex + N*a] - GpMean_pt[J*m + loc + JM*a];
							GpVar_pt[J*m + loc + JM*a] = GpVar_pt[J*m + loc + JM*a] + temp1*temp1/(float)(firmCt_pt[loc] - 1);
						}
					}
				}

				/* Mean/Var of all firm-level means in the market */
				if (totFirm > 0)
				{
					for(j = 0; j < J; j++)
					{
						for (a = 0; a < A; a++)
						{
							GpMean_mean_pt[m + M*a] = GpMean_mean_pt[m + M*a] + GpMean_pt[J*m + j + JM*a]/(float)(totFirm);
							GpVar_mean_pt[m + M*a]  = GpVar_mean_pt[m + M*a]  + GpVar_pt[J*m + j + JM*a]/(float)(totFirm);
						}
					}
				}
			}
		}
	}

	free(rowVec);
	free(firmCt);

}
